import tensorflow as tf
import numpy as np
from supervised_lenet import inference
import matplotlib.pyplot as plt
from foolbox.models import TensorFlowModel
from foolbox.criteria import TargetClass
from foolbox.attacks import CarliniWagnerL2Attack
from tensorflow.contrib.learn.python.learn.datasets.mnist import read_data_sets

mnist = read_data_sets("MNIST_data/",one_hot=True)

images = tf.placeholder(tf.float32,shape=(None,784))
y = tf.placeholder(tf.float32, [None, 10])
loss,logits = inference(images,y,keep_prob=1.0)
saver = tf.train.Saver()

with tf.Session() as session:
    saver.restore(session, tf.train.latest_checkpoint('./supervised_model'))
    print('finish loading model!')

    model = TensorFlowModel(images, logits, bounds=(0, 1))
    criterion = TargetClass(4)
    attack = CarliniWagnerL2Attack(model,criterion)

    # image = mnist.test.images[30]
    # label = np.argmax(model.predictions(image))
    # adversarial = attack(image,label)
    # print("true_label:",np.argmax(mnist.test.labels[30]))
    # print("adversarial_label:",np.argmax(model.predictions(adversarial)))
    per_dataset = np.zeros((4,784))
    for i in range(4):
        image = mnist.test.images[i]
        label = np.argmax(model.predictions(image))
        adversarial = attack(image,label)
        per_dataset[i] = adversarial
        print("true_label:",np.argmax(mnist.test.labels[i]))
        print("adversarial_label:",np.argmax(model.predictions(adversarial)))
        print("making" + str(i) + "picture")
    np.save("supervised_test_picture",per_dataset)

    plt.figure()
    plt.imshow(per_dataset[0].reshape(28,28),cmap='gray')
    plt.xticks([])
    plt.yticks([])
    plt.savefig('ering/0.svg')
    plt.show()

    plt.figure()
    plt.imshow(per_dataset[1].reshape(28,28),cmap='gray')
    plt.xticks([])
    plt.yticks([])
    plt.savefig('ering/1.svg')
    plt.show()

    plt.figure()
    plt.imshow(per_dataset[2].reshape(28,28),cmap='gray')
    plt.xticks([])
    plt.yticks([])
    plt.savefig('ering/2.svg')
    plt.show()

    plt.figure()
    plt.imshow(per_dataset[3].reshape(28, 28), cmap='gray')
    plt.xticks([])
    plt.yticks([])
    plt.savefig('ering/3.svg')
    plt.show()